import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

# 1. 生成数据集
# 生成从 -PI 到 PI 之间的一千个点
x = torch.unsqueeze(torch.linspace(-np.pi, np.pi, 1000), dim=1)
# y 是 x 的 sin 值，并加入一些随机噪声
y = torch.sin(x) + 0.05 * torch.randn(x.size())

# 2. 定义神经网络模型
class SinNet(nn.Module):
    def __init__(self):
        super(SinNet, self).__init__()
        # 定义一个简单的三层全连接网络
        # 输入层 -> 隐藏层1 -> 隐藏层2 -> 输出层
        self.net = nn.Sequential(
            nn.Linear(1, 10),
            nn.ReLU(),
            nn.Linear(10, 10),
            nn.ReLU(),
            nn.Linear(10, 1)
        )

    def forward(self, x):
        return self.net(x)

# 实例化模型
model = SinNet()
print(model)

# 3. 定义损失函数和优化器
# 使用均方误差损失
criterion = nn.MSELoss()
# 使用Adam优化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 4. 训练模型
epochs = 2000 # 训练周期
for epoch in range(epochs):
    # 前向传播
    outputs = model(x)
    loss = criterion(outputs, y)

    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch+1) % 200 == 0:
        print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}')

# 5. 保存模型
# 保存整个模型
torch.save(model, '../output_model/sinx_model.pth')
# 或者只保存模型的状态字典（推荐，更灵活）
torch.save(model.state_dict(), '../output_model/sinx_model_state_dict.pth')

# 为了在C++中使用，我们需要将模型导出为Torch Script格式
# 创建一个模型的实例
scripted_model = torch.jit.script(model)
scripted_model.save('../output_model/sinx_model_scripted.pt')
print("\n模型已导出为 sinx_model_scripted.pt")


# 6. 训练结果可视化
predicted = model(x).detach()
plt.plot(x.numpy(), y.numpy(), 'ro', alpha=0.5, label='Original data')
plt.plot(x.numpy(), predicted.numpy(), 'b', label='Fitted line')
plt.legend()
plt.title('Sine Wave Fitting')
plt.show()